import os

import joblib
import matplotlib
import numpy as np
from xgboost import XGBRegressor

matplotlib.use('Agg')
from sklearn.model_selection import GridSearchCV
from sklearn.multioutput import MultiOutputRegressor
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler


def augment_data(X, Y, n_samples, noise_level):
    X_augmented = [X + noise_level * np.random.randn(*X.shape) for _ in range(n_samples)]
    y_augmented = [Y] * n_samples
    return np.vstack(X_augmented), np.vstack(y_augmented)


def train_and_save_model(X_train, Y_train, X_test, Y_test, model_suffix, model_path,
                         param_grid, cv, scoring):

    pipeline = Pipeline([
        ('scaler', StandardScaler()),
        ('regressor', MultiOutputRegressor(XGBRegressor()))
    ])
    grid_search = GridSearchCV(pipeline, param_grid, cv=cv, scoring=scoring, n_jobs=1)
    grid_search.fit(X_train, Y_train)

    best_model = grid_search.best_estimator_
    best_params = grid_search.best_params_

    os.makedirs(model_path, exist_ok=True)
    joblib.dump(best_model, os.path.join(model_path, f'{model_suffix}_model.pkl'))
    Y_pred = np.clip(np.round(best_model.predict(X_test)), 0, 99).astype(int)

    return Y_pred, best_model, best_params
